"""
Neuron
Represents one neuron in a network
"""

from Constants import *
import random
import math

"""
Represents a spike from another neuron
"""
class Spike:
    def __init__( self ):
        self.source = None # source neuron
        self.arrival = 0.0 # arrival time

"""
Represents a connection to another neuron (synapses in particular - axons are ignored after we convert them to transit times
"""
class Neighbor:
    def __init__( self ):
        self.source = None # source neuron of neighbor
        self.destination = None # destination neuron of neighbor
        self.weight = 0.0 # weight of the connection
        self.ddt_weight = 0.0 # weight of the connection @ previous time step (weights governed by 2nd order diffeq)
        self.distance = 0.0 # distance between neurons
        self.last_time = 0.0 # time of last arrival of a spike via this synapse

class Neuron:
    """
    FIRE ZE MISSILES
    """
    def __init__( self ):
        self.v = 0.0
        self.u = 0.0
        self.w = 0.0
        self.R = 0.0
        self.g_A = 0.0
        self.g_N = 0.0
        self.g_GA = 0.0
        self.g_GB = 0.0

        self.a = 0.0
        self.b = 0.0
        self.c = 0.0
        self.d = 0.0
        self.U = 0.0
        self.F = 0.0
        self.D = 0.0

        self.theta = 0.0
        self.phi = 0.0

        self.far_theta = 0.0
        self.far_phi = 0.0

        self.exc = False

        self.spikes = list()
        self.neighbors = list()
        self.last_fired = 0.0 # Time of last firing
    """ 
    Update all dynamic variables according to smooth ODE's and first- (or second-) order Euler
    """
    def smooth_step( self, time ):
        # Compute the synaptic current and the noisy input current
        I_syn = self.g_A * self.v + self.g_N * self.v * ( ( self.v + 80.0 ) / 60.0 )**2 / ( 1.0 + ( ( self.v + 80.0 ) / 60.0 )**2 ) + self.g_GA * ( self.v + 70.0 ) + self.g_GB * ( self.v + 90.0 )
        
        if time >= 500:
            I_noise = NOISE_MID + NOISE_AMP*math.sin((time/500.0)*2.0*math.pi)
            # I_noise = NOISE_MID*(time-500.0)/(100.0)
        else:
            I_noise = 0.0

        # Smooth updating of internal dynamic variables.
        self.v = self.v + DELTA_T * ( 0.04 * self.v**2 + 5.0 * self.v + 140.0 - self.u - I_syn - I_noise )
        self.u = self.u + DELTA_T * ( self.a * self.b * self.v - self.a * self.u )
        self.R = self.R + DELTA_T * ( ( 1.0 - self.R ) / self.D )
        self.w = self.w + DELTA_T * ( ( self.U - self.w ) / self.F )
        self.g_A = self.g_A * ( 1.0 - DELTA_T / 5.0 )
        self.g_N = self.g_N * ( 1.0 - DELTA_T / 150.0 )
        self.g_GA = self.g_GA * ( 1.0 - DELTA_T / 6.0 )
        self.g_GB = self.g_GB * ( 1.0 - DELTA_T / 150.0 )
        
        # Smooth updating of network weights. Cij is governed by a second-order linear diffeq so this can be transformed into a two-dimensional linear system
        for neighbor in self.neighbors:
            w = neighbor.weight
            ddtw = neighbor.ddt_weight
            
            neighbor.weight = w + ddtw*DELTA_T
            neighbor.ddt_weight = neighbor.ddt_weight + 10.0**(-4.0)*(LEARNING - neighbor.ddt_weight)*DELTA_T
            
            # Clip network weights to [0, 1/2]
            if neighbor.weight <= 0.0:
                neighbor.weight = 0.0
            elif neighbor.weight >= 0.5:
                neighbor.weight = 0.5

    """
    Send a spike to all neighbors
    """
    def send_spike( self, src, time ):

        self.u = self.u + self.d
        self.v = self.c # after-spike resetting of v and u.
        self.last_fired = time

        # Send spikes down the line to connecting neighbors
        for neighbor in self.neighbors:
            if neighbor.source == self: # Only send a spike to other neurons if the firing neuron is the source.
                new = Spike()
                new.source = src
                new.arrival = neighbor.distance
                neighbor.destination.spikes.append( new )
                
        # STPD update: firing a spike strengthens pre-synaptic connections
        for neighbor in self.neighbors:
            if neighbor.destination == self:
                time_diff = time - neighbor.last_time
                time_diff = A_PLUS*math.exp(-time_diff/TAU_PLUS)
                neighbor.ddt_weight = neighbor.ddt_weight + time_diff         
                            
    """
    Spike arrival mechanism
    """
    def receive_spike( self, time ):
        for spk in self.spikes:
            if spk.arrival <= 0: # Spikes are arriving
                # Short-term synaptic plasticity updates
                tempR = -self.R*self.w
                tempw = self.U*(1.0-self.w)
                self.R = self.R + tempR
                self.w = self.w + tempw 
                
                # Synaptic conductance update
                if spk.source.exc == True:
                    for neighbor in spk.source.neighbors:
                        if neighbor.destination == self:
                            self.g_A += neighbor.weight*spk.source.R*spk.source.w
                            # print "Hallelujah! Neighbor.weight " + str(neighbor.weight)
                            self.g_N += neighbor.weight*spk.source.R*spk.source.w

                else:
                    self.g_GA += spk.source.R*spk.source.w
                    self.g_GB += spk.source.R*spk.source.w
                    
                # Update time of last arrival
                for neighbor in self.neighbors:
                    if neighbor.source == spk.source:
                        neighbor.last_time = time
                        
                # Receiving a spike weakens connections (STPD update)
                for neighbor in self.neighbors:
                    time_diff = time - self.last_fired
                    time_diff = A_MINUS*math.exp(-time_diff/TAU_MINUS)
                    neighbor.ddt_weight = neighbor.ddt_weight - time_diff                    
                    
                self.spikes.remove(spk)
                
            else:
                spk.arrival = spk.arrival - 1 # Count down!  Then shalt thou count to three, no more, no less.
    """
    Return the interesting data values of the neuron
    """
    def get_vals( self, time ):
        li = [ self.v, self.u, self.R, self.w, self.g_A, self.g_GA ]
        return li
